'''
Created on Jul 3, 2011

@author: Nam Khanh
'''
from dep import *
from liblinear import *
from liblinearutil import *

n_gram = 3

def to_node(line):
    """
    """
    token = line.split('\t')
    idx = int(token[0])
    form = token[1]
    pos = token[4]
    
    #if training file
    if len(token) < 9:
        head_idx = int(token[6])
        return DepNode(idx = idx, form = form, pos = pos, head_idx = head_idx)
    else:
        target_head = int(token[8])
        return DepNode(idx = idx, form = form, pos = pos, target_head = target_head)

def get_feature(tree, word, history, symbol_table):
    """
    """
    feature = list()
    for i in range(0,n_gram):
        if word-i >= 0:
            f = 'wf_w-' + str(i) + '=' + tree.get_form(word-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
            f = 'pos_w-' + str(i) + '=' + tree.get_pos(word-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
        if history-i >= 0:
            f = 'wf_h-' + str(i) + '=' + tree.get_form(history-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
            f = 'pos_h-' + str(i) + '=' + tree.get_pos(history-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            
    for i in range(0,n_gram):
        if word+i < len(tree):
            f = 'wf_w+' + str(i) + '=' + tree.get_form(word+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
            f = 'pos_w+' + str(i) + '=' + tree.get_pos(word+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
        if history+i < len(tree):
            f = 'wf_h+' + str(i) + '=' + tree.get_form(history+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
                
            f = 'pos_h+' + str(i) + '=' +  tree.get_pos(history+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
       
    f_dict = dict()
    for f in feature:
        f_dict[f] = 1
    
    return f_dict

def predict(sentence, symbol_table, model):
    """
    """
    correct = 0
    tree = DepTree([])
    root_node = DepNode(0, '#$ROOT$#', '#$ROOT$#', -1)
    tree.add(root_node)

    for line in sentence:
        tree.add(to_node(line))
        
    head = []
    for i in range(0, len(tree)):
        head.append(-1)
    
    for word in range(1, len(tree)):
        for i in range(1, word+1):
            history = word - i
            if (head[word] != -1 and head[history] != -1):
                continue
            
            f = get_feature(tree, word, history, symbol_table)
            x0, max_idx = gen_feature_nodearray(f)
            label = liblinear.predict(model, x0)
            
#            print label
            
            if head[word] == -1 and label == 1:
                head[word] = history
            if head[history] == -1 and label == 2:
                head[history] = word                                                                                    

    #print head

    for w in range(1, len(tree)):
        if head[w] == tree.get_target_head(w):
            correct = correct + 1                        

#    print correct 
               
    return correct            

def do_predict(filename):
    """
    """
    print 'Testing n_gram = ', n_gram, '...'
    correct = 0
    
    sentence = []
    model = load_model('english_' + str(n_gram) + '.model')

    symbol_table = dict()
    with open('index_' + str(n_gram) + '.table', 'r') as fin:
        for line in fin.readlines():
            temp = line.split('\t')
            symbol_table[temp[0]] = int(temp[1])
    
    number_sentence = 0
    total = 0
    
    with open(filename, 'r') as fin:
        for line in fin.readlines():
            line = line.strip()
            if len(line) == 0:
                number_sentence = number_sentence + 1
                print "Predicting the sentence ", number_sentence                
                correct = correct + predict(sentence, symbol_table, model)
                total = total + len(sentence)
                print correct, ' vs. ', total
                sentence = []
            else:
                sentence.append(line)                                    
    if len(sentence) > 0:
        number_sentence = number_sentence + 1
        print "Predicting the sentence ", number_sentence                        
        correct = correct + predict(sentence, symbol_table, model)  
        total = total + len(sentence)
    
    print "Total correct = ", correct
    print "UAS = ", correct * 100.0 / total
    
if __name__=='__main__':
    if len(sys.argv) < 3:
        print '\nUsage: python {0} filename -n_gram'.format(sys.argv[0])
        sys.exit(1)
        
    n_gram = int(sys.argv[2])
    do_predict(sys.argv[1])